PyTorch 使用 GPU 进行训练

安装 CUDA

先检查当前环境上是否安装了 CUDA

python -c "import torch; print(torch.cuda.is_available())"

如果没有安装，则到官网找个支持 CUDA 的版本进行安装

https://pytorch.org/get-started/locally/

检查本地的 CUDA 版本，在cmd窗口的命令行执行 nvcc -V 或 nvcc --version 即可查看CUDA.

nvcc --version

可以看到使用的是 12.1 版本的 CUDA，所以我们需要安装 cu121 版本的 PyTorch，这个版本的 PyTorch 适用于 CUDA 11.1、11.3、11.4、12.0 和 12.1。

最终给出的安装地址

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

如果使用的是 poetry 则直接，找到指定依赖就行了，注意下面的 cp311 是 python3.11 的版本号，如果是其他版本的 python 则需要修改。

[tool.poetry.dependencies]
torch = { url = "https://download.pytorch.org/whl/cu121/torch-2.3.0%2Bcu121-cp311-cp311-win_amd64.whl" }
torchvision = { url = "https://download.pytorch.org/whl/cu121/torchvision-0.18.0%2Bcu121-cp311-cp311-win_amd64.whl" }
# ...

如果报错

ERROR: Could not install packages due to an OSError: [WinError 5] 拒绝访问。

添加 --user 选项赋予权限就可以搞定。

pip3 install torch torchvision torchaudio --user --index-url https://download.pytorch.org/whl/cu121

注意：如果是在虚拟目录下则需要修改

到虚拟目录的 pyvenv.cfg 文件夹下，设置 include-system-site-packages 为 true

在 PyTorch 中使用

在 PyTorch 中，使用 GPU 主要涉及将模型和数据移动到 GPU 上。以下是如何使用 GPU 的基本步骤和示例：

检查 GPU 的可用性:
```
use_cuda = torch.cuda.is_available()
```

定义设备:

device = torch.device("cuda:0" if use_cuda else "cpu")

将模型移动到 GPU:
```
model = model.to(device)
```

将数据移动到 GPU:

inputs, labels = inputs.to(device), labels.to(device)

以下是一个简单的示例，展示如何在 PyTorch 中使用 GPU：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

# 检查 GPU 的可用性
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")

# 创建模型并移动到 GPU
model = SimpleModel().to(device)

# 创建一个优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 生成一些随机数据
data = torch.randn(5, 10).to(device)
target = torch.randn(5, 10).to(device)

# 训练模型
optimizer.zero_grad()
output = model(data)
loss = nn.MSELoss()(output, target)
loss.backward()
optimizer.step()

print("Loss:", loss.item())

注意：在上述代码中，我们使用了 .to(device) 方法将模型和数据移动到 GPU（如果可用）。这是一个简洁的方法，可以确保代码在没有 GPU 的机器上也能运行，只是运行速度会慢一些。

使用 GPU 并行训练

使用 PyTorch 进行 GPU 并行训练有多种方法，下面是一些常见的方法：

使用单个 GPU: 你只需要确保你的模型和数据都在 GPU 上。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
data, target = data.to(device), target.to(device)

使用 DataParallel 进行多 GPU 并行训练: 这是一个简单的并行化方法，它会将输入数据拆分到多个 GPU，然后每个 GPU 计算其部分的前向和后向传递，最后在主 GPU 上合并结果。
```
if torch.cuda.device_count() > 1:
    print("Using", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

model = model.to(device)
```

使用 DistributedDataParallel 进行多 GPU 并行训练: 如果你需要在多个节点上进行分布式训练或者使用更复杂的并行策略，DistributedDataParallel 是一个好的选择。

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# 初始化过程 (例如，使用 nccl 通信后端)
dist.init_process_group(backend='nccl')

model = model.to(device)
if torch.cuda.device_count() > 1:
    print("Using", torch.cuda.device_count(), "GPUs!")
    model = DistributedDataParallel(model, device_ids=[device])

与此同时，使用 DistributedSampler 可以确保数据集在多个 GPU 之间正确地分布。

from torch.utils.data.distributed import DistributedSampler

train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)

运行分布式训练通常需要使用 torch.distributed.launch 帮助启动脚本。

自定义 GPU 并行策略: 除了上述工具外，你还可以直接使用 PyTorch 提供的低级 CUDA 功能（如 torch.cuda.comm.scatter 和 torch.cuda.comm.gather）来实现自定义的并行策略。
注意事项:
- 当使用多 GPU 并行训练时，由于每个 GPU 都会处理批次数据的一个子集，因此你可能需要调整学习率或其他超参数。
- 确保你的数据加载和批次大小可以适应多 GPU 训练。不同的 GPU 之间可能会存在通信开销，因此确保你的模型和批次大小是适合多 GPU 的。

使用这些策略，你可以有效地使用 PyTorch 在多个 GPU 上并行训练模型。

安装 CUDA​

在 PyTorch 中使用​

使用 GPU 并行训练​

安装 CUDA

在 PyTorch 中使用

使用 GPU 并行训练